Simple image classification, vertical versus horizontal stripes ===============================================

# pylint: disable=invalid-name

from time import time as tictoc

import numpy as np

from qtealeaves.convergence_parameters import TNConvergenceParameters
from qtealeaves.emulator import MPS
from qtealeaves.mpos.mldatampo import MLDataMPO
from qtealeaves.tensors import TensorBackend


def generate_stripe_image(image_size, conv_params, tensor_backend, vertical=True):
    """
    Generate an image with stripes in MPS form

    Parameters
    ----------
    image_size : tuple
        Shape of the image
    conv_params : TNConvergenceParameters
        Convergence parameters of the MPS
    vertical : bool, optional
        Direction of the stripes, by default True

    Returns
    -------
    MPS
        MPS representation of the image
    """
    image = np.zeros(image_size)

    if vertical:
        image[:, np.random.randint(0, image_size[1])] = 1
    else:
        image[np.random.randint(0, image_size[0]), :] = 1
    image += np.random.normal(0, 0.1, np.prod(image_size)).reshape(image_size)
    tensorlist = [
        # pylint: disable-next=too-many-function-args
        np.array([np.cos(pixel), np.sin(pixel)]).reshape(1, 2, 1)
        for pixel in image.reshape(-1)
    ]
    mps_image = MPS.from_tensor_list(
        tensorlist,
        conv_params=conv_params,
        tensor_backend=tensor_backend,
    )
    mps_image.right_canonize(0)
    mps_image.normalize()

    return mps_image


def get_learning_rate_func(num_sweeps, final_rate):
    """Learning rate, exponentially decaying"""
    coeff = np.log(final_rate) / (1 - num_sweeps)

    def learning_rate(x_value, coeff=coeff):
        return np.exp(-coeff * x_value)

    return learning_rate


# pylint: disable-next=too-many-locals
def main(batch_size=19, ansatz=MPS, image_size=5, train_size=20, test_size=10):
    """Main method for classification problem via tensor network machine learning."""
    tic = tictoc()
    np.random.seed([11, 13, 17, 19])
    image_size_ = (image_size, image_size)
    num_sweeps = 20
    conv_parameters = TNConvergenceParameters(
        max_bond_dimension=8, data_type="S", max_iter=num_sweeps
    )
    learning_rate = get_learning_rate_func(num_sweeps, 1e-5)

    tensor_backend = TensorBackend(dtype=np.float32)
    dataset = [
        generate_stripe_image(image_size_, conv_parameters, tensor_backend, ii % 2)
        for ii in range(train_size)
    ]
    labels = np.array([ii % 2 for ii in range(train_size)])

    ml_data_mpo = MLDataMPO(dataset, labels, batch_size, tensor_backend)

    classificator = ansatz.ml_initial_guess(
        conv_parameters, tensor_backend, "superposition-data", ml_data_mpo, dataset
    )

    _, _ = classificator.ml_optimize(ml_data_mpo, learning_rate, "linkfree")

    predictions = classificator.ml_predict(dataset, 6)
    print(f"Train accuracy: {np.sum(1-np.abs(labels-predictions))/train_size}")

    test_dataset = [
        generate_stripe_image(image_size_, conv_parameters, tensor_backend, ii % 2)
        for ii in range(test_size)
    ]
    test_labels = np.array([ii % 2 for ii in range(test_size)])

    predictions = classificator.ml_predict(test_dataset, 1)

    print(f"Test accuracy: {np.sum(1-np.abs(test_labels-predictions))/test_size}")

    time = tictoc() - tic
    print(
        f"\nExample `{__file__}` ran successfully in {time}s; "
        + "no asserts implemented."
    )


if __name__ == "__main__":
    main()